import torch.nn as nn
from . import block as B
import torch

# For any upscale factors
class IMDN_AS(nn.Module):
    def __init__(self, in_nc=3, nf=64, num_modules=6, out_nc=3, upscale=4):
        super(IMDN_AS, self).__init__()

        self.fea_conv = nn.Sequential(B.conv_layer(in_nc, nf, kernel_size=3, stride=2),
                                      nn.LeakyReLU(0.05),
                                      B.conv_layer(nf, nf, kernel_size=3, stride=2))

        # IMDBs
        self.IMDB1 = B.IMDModule(in_channels=nf)
        self.IMDB2 = B.IMDModule(in_channels=nf)
        self.IMDB3 = B.IMDModule(in_channels=nf)
        self.IMDB4 = B.IMDModule(in_channels=nf)
        self.IMDB5 = B.IMDModule(in_channels=nf)
        self.IMDB6 = B.IMDModule(in_channels=nf)
        self.c = B.conv_block(nf * num_modules, nf, kernel_size=1, act_type='lrelu')

        self.LR_conv = B.conv_layer(nf, nf, kernel_size=3)

        upsample_block = B.pixelshuffle_block
        self.upsampler = upsample_block(nf, out_nc, upscale_factor=upscale)


    def forward(self, input):
        out_fea = self.fea_conv(input)
        out_B1 = self.IMDB1(out_fea)
        out_B2 = self.IMDB2(out_B1)
        out_B3 = self.IMDB3(out_B2)
        out_B4 = self.IMDB4(out_B3)
        out_B5 = self.IMDB5(out_B4)
        out_B6 = self.IMDB6(out_B5)

        out_B = self.c(torch.cat([out_B1, out_B2, out_B3, out_B4, out_B5, out_B6], dim=1))
        out_lr = self.LR_conv(out_B) + out_fea
        output = self.upsampler(out_lr)
        return output

class IMDN(nn.Module):
    def __init__(self, in_nc=3, nf=64, num_modules=6, out_nc=3, upscale=4):
        super(IMDN, self).__init__()

        self.fea_conv = B.conv_layer(in_nc, nf, kernel_size=3)

        # IMDBs
        self.IMDB1 = B.IMDModule(in_channels=nf)
        self.IMDB2 = B.IMDModule(in_channels=nf)
        self.IMDB3 = B.IMDModule(in_channels=nf)
        self.IMDB4 = B.IMDModule(in_channels=nf)
        self.IMDB5 = B.IMDModule(in_channels=nf)
        self.IMDB6 = B.IMDModule(in_channels=nf)
        self.IMDB7 = B.IMDModule(in_channels=nf)
        self.IMDB8 = B.IMDModule(in_channels=nf)
        #self.c = B.conv_block(nf * num_modules, nf, kernel_size=1, act_type='lrelu')

        self.LR_conv = B.conv_layer(nf, nf, kernel_size=3)

        upsample_block = B.pixelshuffle_block
        #upsample_block2 = B.pixelshuffle_block

        #self.upsampler0 = upsample_block2(nf, nf//4, upscale_factor=upscale // 2)

        self.upsampler = upsample_block(nf, out_nc, upscale_factor=upscale)

       # s#elf.downsample  =  B.PixelUnshuffle(2)
#

        self.fusion1 =  B.conv_layer(nf*2, nf, kernel_size=1)
        self.fusion2 = B.conv_layer(nf * 2, nf, kernel_size=1)
        self.fusion3 = B.conv_layer(nf * 2, nf, kernel_size=1)

        self.fusion4 = B.conv_layer(nf * 2, nf, kernel_size=1)
        self.fusion5 = B.conv_layer(nf * 2, nf, kernel_size=1)
        self.fusion6 = B.conv_layer(nf * 2, nf, kernel_size=1)

        self.fusion7 = B.conv_layer(nf * 2, nf, kernel_size=1)



    def forward(self, input):
        out_fea = self.fea_conv(input)
        out_B1 = self.IMDB1(out_fea)
        out_B2 = self.IMDB2(out_B1)
        #out_B2 = self.upsampler0(out_B2)
        out_B3 = self.IMDB3(out_B2)
        out_B4 = self.IMDB4(out_B3)
        out_B5 = self.IMDB5(out_B4)
        out_B6 = self.IMDB6(out_B5)
        out_B7 = self.IMDB7(out_B6)
        out_B8 = self.IMDB8(out_B7)


        F1 = self.fusion1 (torch.cat([out_B1, out_B2],1))
        F2 = self.fusion2 (torch.cat([out_B3, out_B4],1))
        F3 = self.fusion3 (torch.cat([out_B5, out_B6], 1))
        F4 = self.fusion4 (torch.cat([out_B7, out_B8], 1))

        F5 = self.fusion5 (torch.cat([F1, F2], 1))
        F6 = self.fusion6(torch.cat([F3, F4], 1))

        F7 = self.fusion7(torch.cat([F5,F6], 1))





        #out_B = self.c(torch.cat([out_B1, out_B2, out_B3, out_B4, out_B5, out_B6], dim=1))
        out_lr = self.LR_conv(F7) + out_fea #self.LR_conv(out_B) + out_fea
        output = self.upsampler(out_lr)
        return output

class IMDN_RTC(nn.Module):
    def __init__(self, in_nc=3, nf=12, num_modules=5, out_nc=3, upscale=2):
        super(IMDN_RTC, self).__init__()

        fea_conv = [B.conv_layer(in_nc, nf, kernel_size=3)]
        rb_blocks = [B.IMDModule_speed(in_channels=nf) for _ in range(num_modules)]
        LR_conv = B.conv_layer(nf, nf, kernel_size=1)

        upsample_block = B.pixelshuffle_block
        upsampler = upsample_block(nf, out_nc, upscale_factor=upscale)

        self.model = B.sequential(*fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv)),
                                  *upsampler)

    def forward(self, input):
        output = self.model(input)
        return output
